





<!DOCTYPE html>
<html class="writer-html5" lang="zh-CN" >
<head>
  <meta charset="utf-8">
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>How to optimize convolution using TensorCores &mdash; tvm 0.8.dev1982 文档</title>
  

  
  <link rel="stylesheet" href="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/css/bootstrap.min.css" integrity="sha384-Gn5384xqQ1aoWXA+058RXPxPg6fy4IWvTNh0E263XmFcJlSAwiGgFAW/dAiS6JXm" crossorigin="anonymous">
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/gallery.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="../../_static/css/tlcpack_theme.css" type="text/css" />

  
  
    <link rel="shortcut icon" href="../../_static/tvm-logo-square.png"/>
  

  
  
  
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
        <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
        <script src="../../_static/jquery.js"></script>
        <script src="../../_static/underscore.js"></script>
        <script src="../../_static/doctools.js"></script>
        <script src="../../_static/translations.js"></script>
    
    <script type="text/javascript" src="../../_static/js/theme.js"></script>

    
    <script type="text/javascript" src="../../_static/js/tlcpack_theme.js"></script>
    <link rel="index" title="索引" href="../../genindex.html" />
    <link rel="search" title="搜索" href="../../search.html" />
    <link rel="next" title="Auto-Tune with Templates and AutoTVM" href="../tune_with_autotvm/index.html" />
    <link rel="prev" title="如何在GPU上优化卷积" href="opt_conv_cuda.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    
<header class="header">
    <div class="innercontainer">
      <div class="headerInner d-flex justify-content-between align-items-center">
          <div class="headerLogo">
               <a href="https://tvm.apache.org/"><img src=https://tvm.apache.org/assets/images/logo.svg alt="logo"></a>
          </div>

          <div id="headMenu" class="headerNav">
            <button type="button" id="closeHeadMenu" class="navCloseBtn"><img src="../../_static/img/close-icon.svg" alt="Close"></button>
             <ul class="nav">
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/community>Community</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/download>Download</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/vta>VTA</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/blog>Blog</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvm.apache.org/docs>Docs</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmconf.org>Conference</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://github.com/apache/tvm/>Github</a>
                </li>
                <li class="nav-item">
                   <a class="nav-link" href=https://tvmchinese.github.io/declaration_zh_CN.html>About-Translators</a>
                </li>
             </ul>
               <div class="responsivetlcdropdown">
                 <button type="button" class="btn-link">
                   ASF
                 </button>
                 <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                 </ul>
               </div>
          </div>
            <div class="responsiveMenuIcon">
              <button type="button" id="menuBtn" class="btn-menu"><img src="../../_static/img/menu-icon.svg" alt="Menu Icon"></button>
            </div>

            <div class="tlcDropdown">
              <div class="dropdown">
                <button type="button" class="btn-link dropdown-toggle" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">
                  ASF
                </button>
                <div class="dropdown-menu dropdown-menu-right">
                  <ul>
                     <li>
                       <a href=https://apache.org/>Apache Homepage</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/licenses/>License</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/sponsorship.html>Sponsorship</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/security/>Security</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/foundation/thanks.html>Thanks</a>
                     </li>
                     <li>
                       <a href=https://www.apache.org/events/current-event>Events</a>
                     </li>
                     <li>
                       <a href=https://www.zhihu.com/column/c_1429578595417563136>Zhihu</a>
                     </li>
                  </ul>
                </div>
              </div>
          </div>
       </div>
    </div>
 </header>
 
    <nav data-toggle="wy-nav-shift" class="wy-nav-side fixed">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../index.html">
          

          
            
            <img src="../../_static/tvm-logo-small.png" class="logo" alt="Logo"/>
          
          </a>

          
            
            
                <div class="version">
                  0.8.dev1982
                </div>
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption" role="heading"><span class="caption-text">如何开始</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../install/index.html">安装 TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../contribute/index.html">贡献者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">用户引导</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="../../tutorial/index.html">User Tutorial</a></li>
<li class="toctree-l1 current"><a class="reference internal" href="../index.html">How To Guides</a><ul class="current">
<li class="toctree-l2"><a class="reference internal" href="../compile_models/index.html">编译深度学习模型</a></li>
<li class="toctree-l2"><a class="reference internal" href="../deploy/index.html">TVM 部署模型和集成</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_relay/index.html">Work With Relay</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_schedules/index.html">Work With Tensor Expression and Schedules</a></li>
<li class="toctree-l2 current"><a class="reference internal" href="index.html">优化张量算子</a><ul class="current">
<li class="toctree-l3"><a class="reference internal" href="opt_gemm.html">如何在CPU上优化GEMM（通用矩阵乘）</a></li>
<li class="toctree-l3"><a class="reference internal" href="opt_conv_cuda.html">如何在GPU上优化卷积</a></li>
<li class="toctree-l3 current"><a class="current reference internal" href="#">How to optimize convolution using TensorCores</a><ul>
<li class="toctree-l4"><a class="reference internal" href="#tensorcore-introduction">TensorCore Introduction</a></li>
<li class="toctree-l4"><a class="reference internal" href="#preparation-and-algorithm">准备和算法</a></li>
<li class="toctree-l4"><a class="reference internal" href="#memory-scope">Memory Scope</a></li>
<li class="toctree-l4"><a class="reference internal" href="#define-tensor-intrinsic">Define Tensor Intrinsic</a></li>
<li class="toctree-l4"><a class="reference internal" href="#scheduling-the-computation">Scheduling the Computation</a></li>
<li class="toctree-l4"><a class="reference internal" href="#lowering-computation-to-intrinsics">Lowering Computation to Intrinsics</a></li>
<li class="toctree-l4"><a class="reference internal" href="#generate-cuda-kernel">Generate CUDA Kernel</a></li>
<li class="toctree-l4"><a class="reference internal" href="#summary">总结</a></li>
</ul>
</li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="../tune_with_autotvm/index.html">Auto-Tune with Templates and AutoTVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../tune_with_autoscheduler/index.html">Use AutoScheduler for Template-Free Scheduling</a></li>
<li class="toctree-l2"><a class="reference internal" href="../work_with_microtvm/index.html">Work With microTVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../extend_tvm/index.html">Extend TVM</a></li>
<li class="toctree-l2"><a class="reference internal" href="../profile/index.html">Profile Models</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../errors.html">Handle TVM Errors</a></li>
<li class="toctree-l2"><a class="reference internal" href="../../faq.html">常见提问</a></li>
</ul>
</li>
</ul>
<p class="caption" role="heading"><span class="caption-text">开发者引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../dev/tutorial/index.html">Developer Tutorial</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../dev/how_to/how_to.html">开发者指南</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">架构指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../arch/index.html">Design and Architecture</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">主题引导</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../topic/microtvm/index.html">microTVM：裸机使用TVM</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../topic/vta/index.html">VTA: Versatile Tensor Accelerator</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">参考指南</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../reference/langref/index.html">语言参考</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/python/index.html">Python API</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/api/links.html">Other APIs</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../reference/publications.html">Publications</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../genindex.html">索引</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
      
      <nav class="wy-nav-top" aria-label="top navigation" data-toggle="wy-nav-top">
        
            <div class="togglemenu">

            </div>
            <div class="nav-content">
              <!-- tvm -->
              Table of content
            </div>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        

          




















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../index.html">Docs</a> <span class="br-arrow">></span></li>
        
          <li><a href="../index.html">How To Guides</a> <span class="br-arrow">></span></li>
        
          <li><a href="index.html">优化张量算子</a> <span class="br-arrow">></span></li>
        
      <li>How to optimize convolution using TensorCores</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
            
            <a href="../../_sources/how_to/optimize_operators/opt_conv_tensorcore.rst.txt" rel="nofollow"> <img src="../../_static//img/source.svg" alt="viewsource"/></a>
          
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <div class="sphx-glr-download-link-note admonition note">
<p class="admonition-title">注解</p>
<p>Click <a class="reference internal" href="#sphx-glr-download-how-to-optimize-operators-opt-conv-tensorcore-py"><span class="std std-ref">here</span></a> to download the full example code</p>
</div>
<div class="sphx-glr-example-title section" id="how-to-optimize-convolution-using-tensorcores">
<span id="opt-conv-tensorcore"></span><span id="sphx-glr-how-to-optimize-operators-opt-conv-tensorcore-py"></span><h1>How to optimize convolution using TensorCores<a class="headerlink" href="#how-to-optimize-convolution-using-tensorcores" title="永久链接至标题">¶</a></h1>
<p><strong>作者</strong>: <a class="reference external" href="https://github.com/Hzfengsy">Siyuan Feng</a></p>
<p>In this tutorial, we will demonstrate how to write a high performance convolution
schedule using TensorCores in TVM. In this example, we assume the input to
convolution has a large batch. We strongly recommend covering the <a class="reference internal" href="opt_conv_cuda.html#opt-conv-gpu"><span class="std std-ref">如何在GPU上优化卷积</span></a> tutorial first.</p>
<div class="section" id="tensorcore-introduction">
<h2>TensorCore Introduction<a class="headerlink" href="#tensorcore-introduction" title="永久链接至标题">¶</a></h2>
<p>Each Tensor Core provides a 4x4x4 matrix processing array that operates
<code class="code docutils literal notranslate"><span class="pre">D</span> <span class="pre">=</span> <span class="pre">A</span> <span class="pre">*</span> <span class="pre">B</span> <span class="pre">+</span> <span class="pre">C</span></code>, where A, B, C and D are 4x4 matrices as Figure shows.
The matrix multiplication inputs A and B are FP16 matrices, while the accumulation
matrices C and D may be FP16 or FP32 matrices.</p>
<p>However, CUDA programmers can only use warp-level primitive
<code class="code docutils literal notranslate"><span class="pre">wmma::mma_sync(acc_frag,</span> <span class="pre">a_frag,</span> <span class="pre">b_frag,</span> <span class="pre">acc_frag)</span></code> to perform
16x16x16 half-precision matrix multiplication on tensor cores. Before invoking
the matrix multiplication, programmers must load data from memory into registers
with primitive <code class="code docutils literal notranslate"><span class="pre">wmma::load_matrix_sync</span></code>, explicitly. The NVCC compiler translates
that primitive into multiple memory load instructions. At run time, every thread loads
16 elements from matrix A and 16 elements from B.</p>
</div>
<div class="section" id="preparation-and-algorithm">
<h2>准备和算法<a class="headerlink" href="#preparation-and-algorithm" title="永久链接至标题">¶</a></h2>
<p>We use the fixed size for input tensors with 256 channels and 14 x 14 dimensions.
The batch size is 256. Convolution filters contain 512 filters of size 3 x 3.
We use stride size 1 and padding size 1 for the convolution. In the example, we use
NHWCnc memory layout.The following code defines the convolution algorithm in TVM.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span> <span class="nn">tvm</span>
<span class="kn">from</span> <span class="nn">tvm</span> <span class="k">import</span> <span class="n">te</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">tvm.contrib</span> <span class="k">import</span> <span class="n">nvcc</span>

<span class="c1"># The sizes of inputs and filters</span>
<span class="n">batch_size</span> <span class="o">=</span> <span class="mi">256</span>
<span class="n">height</span> <span class="o">=</span> <span class="mi">14</span>
<span class="n">width</span> <span class="o">=</span> <span class="mi">14</span>
<span class="n">in_channels</span> <span class="o">=</span> <span class="mi">256</span>
<span class="n">out_channels</span> <span class="o">=</span> <span class="mi">512</span>
<span class="n">kernel_h</span> <span class="o">=</span> <span class="mi">3</span>
<span class="n">kernel_w</span> <span class="o">=</span> <span class="mi">3</span>
<span class="n">pad_h</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">pad_w</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">stride_h</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">stride_w</span> <span class="o">=</span> <span class="mi">1</span>

<span class="c1"># TensorCore shape</span>
<span class="n">block_size</span> <span class="o">=</span> <span class="mi">16</span>

<span class="k">assert</span> <span class="n">batch_size</span> <span class="o">%</span> <span class="n">block_size</span> <span class="o">==</span> <span class="mi">0</span>
<span class="k">assert</span> <span class="n">in_channels</span> <span class="o">%</span> <span class="n">block_size</span> <span class="o">==</span> <span class="mi">0</span>
<span class="k">assert</span> <span class="n">out_channels</span> <span class="o">%</span> <span class="n">block_size</span> <span class="o">==</span> <span class="mi">0</span>

<span class="c1"># Input feature map: (N, H, W, IC, n, ic)</span>
<span class="n">data_shape</span> <span class="o">=</span> <span class="p">(</span>
    <span class="n">batch_size</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
    <span class="n">height</span><span class="p">,</span>
    <span class="n">width</span><span class="p">,</span>
    <span class="n">in_channels</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
    <span class="n">block_size</span><span class="p">,</span>
    <span class="n">block_size</span><span class="p">,</span>
<span class="p">)</span>
<span class="c1"># Kernel: (H, W, IC, OC, ic, oc)</span>
<span class="n">kernel_shape</span> <span class="o">=</span> <span class="p">(</span>
    <span class="n">kernel_h</span><span class="p">,</span>
    <span class="n">kernel_w</span><span class="p">,</span>
    <span class="n">in_channels</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
    <span class="n">out_channels</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
    <span class="n">block_size</span><span class="p">,</span>
    <span class="n">block_size</span><span class="p">,</span>
<span class="p">)</span>
<span class="c1"># Output feature map: (N, H, W, OC, n, oc)</span>
<span class="n">output_shape</span> <span class="o">=</span> <span class="p">(</span>
    <span class="n">batch_size</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
    <span class="n">height</span><span class="p">,</span>
    <span class="n">width</span><span class="p">,</span>
    <span class="n">out_channels</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
    <span class="n">block_size</span><span class="p">,</span>
    <span class="n">block_size</span><span class="p">,</span>
<span class="p">)</span>

<span class="c1"># Reduction axes</span>
<span class="n">kh</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">kernel_h</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;kh&quot;</span><span class="p">)</span>
<span class="n">kw</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">kernel_w</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;kw&quot;</span><span class="p">)</span>
<span class="n">ic</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">in_channels</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;ic&quot;</span><span class="p">)</span>
<span class="n">ii</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">block_size</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;ii&quot;</span><span class="p">)</span>

<span class="c1"># Algorithm</span>
<span class="n">A</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">(</span><span class="n">data_shape</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;A&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;float16&quot;</span><span class="p">)</span>
<span class="n">W</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">(</span><span class="n">kernel_shape</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;W&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;float16&quot;</span><span class="p">)</span>
<span class="n">Apad</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">(</span>
    <span class="p">(</span>
        <span class="n">batch_size</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
        <span class="n">height</span> <span class="o">+</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">pad_h</span><span class="p">,</span>
        <span class="n">width</span> <span class="o">+</span> <span class="mi">2</span> <span class="o">*</span> <span class="n">pad_w</span><span class="p">,</span>
        <span class="n">in_channels</span> <span class="o">//</span> <span class="n">block_size</span><span class="p">,</span>
        <span class="n">block_size</span><span class="p">,</span>
        <span class="n">block_size</span><span class="p">,</span>
    <span class="p">),</span>
    <span class="k">lambda</span> <span class="n">n</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">nn</span><span class="p">,</span> <span class="n">ii</span><span class="p">:</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">if_then_else</span><span class="p">(</span>
        <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">all</span><span class="p">(</span><span class="n">h</span> <span class="o">&gt;=</span> <span class="n">pad_h</span><span class="p">,</span> <span class="n">h</span> <span class="o">-</span> <span class="n">pad_h</span> <span class="o">&lt;</span> <span class="n">height</span><span class="p">,</span> <span class="n">w</span> <span class="o">&gt;=</span> <span class="n">pad_w</span><span class="p">,</span> <span class="n">w</span> <span class="o">-</span> <span class="n">pad_w</span> <span class="o">&lt;</span> <span class="n">width</span><span class="p">),</span>
        <span class="n">A</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">h</span> <span class="o">-</span> <span class="n">pad_h</span><span class="p">,</span> <span class="n">w</span> <span class="o">-</span> <span class="n">pad_w</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">nn</span><span class="p">,</span> <span class="n">ii</span><span class="p">],</span>
        <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">const</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="s2">&quot;float16&quot;</span><span class="p">),</span>
    <span class="p">),</span>
    <span class="n">name</span><span class="o">=</span><span class="s2">&quot;Apad&quot;</span><span class="p">,</span>
<span class="p">)</span>
<span class="n">Conv</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">(</span>
    <span class="n">output_shape</span><span class="p">,</span>
    <span class="k">lambda</span> <span class="n">n</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">o</span><span class="p">,</span> <span class="n">nn</span><span class="p">,</span> <span class="n">oo</span><span class="p">:</span> <span class="n">te</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span>
        <span class="n">Apad</span><span class="p">[</span><span class="n">n</span><span class="p">,</span> <span class="n">h</span> <span class="o">*</span> <span class="n">stride_h</span> <span class="o">+</span> <span class="n">kh</span><span class="p">,</span> <span class="n">w</span> <span class="o">*</span> <span class="n">stride_w</span> <span class="o">+</span> <span class="n">kw</span><span class="p">,</span> <span class="n">ic</span><span class="p">,</span> <span class="n">nn</span><span class="p">,</span> <span class="n">ii</span><span class="p">]</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s2">&quot;float32&quot;</span><span class="p">)</span>
        <span class="o">*</span> <span class="n">W</span><span class="p">[</span><span class="n">kh</span><span class="p">,</span> <span class="n">kw</span><span class="p">,</span> <span class="n">ic</span><span class="p">,</span> <span class="n">o</span><span class="p">,</span> <span class="n">ii</span><span class="p">,</span> <span class="n">oo</span><span class="p">]</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s2">&quot;float32&quot;</span><span class="p">),</span>
        <span class="n">axis</span><span class="o">=</span><span class="p">[</span><span class="n">ic</span><span class="p">,</span> <span class="n">kh</span><span class="p">,</span> <span class="n">kw</span><span class="p">,</span> <span class="n">ii</span><span class="p">],</span>
    <span class="p">),</span>
    <span class="n">name</span><span class="o">=</span><span class="s2">&quot;Conv&quot;</span><span class="p">,</span>
<span class="p">)</span>

<span class="n">s</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">create_schedule</span><span class="p">(</span><span class="n">Conv</span><span class="o">.</span><span class="n">op</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">Apad</span><span class="p">]</span><span class="o">.</span><span class="n">compute_inline</span><span class="p">()</span>
</pre></div>
</div>
</div>
<div class="section" id="memory-scope">
<h2>Memory Scope<a class="headerlink" href="#memory-scope" title="永久链接至标题">¶</a></h2>
<p>In traditional GPU schedule, we have global, shared and local memory scope.
To support TensorCores, we add another three special memory scope: <code class="code docutils literal notranslate"><span class="pre">wmma.matrix_a</span></code>,
<code class="code docutils literal notranslate"><span class="pre">wmma.matrix_b</span></code> and <code class="code docutils literal notranslate"><span class="pre">wmma.accumulator</span></code>. On hardware, all fragments scope
stores at the on-chip registers level, the same place with local memory.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Designate the memory hierarchy</span>
<span class="n">AS</span> <span class="o">=</span> <span class="n">s</span><span class="o">.</span><span class="n">cache_read</span><span class="p">(</span><span class="n">Apad</span><span class="p">,</span> <span class="s2">&quot;shared&quot;</span><span class="p">,</span> <span class="p">[</span><span class="n">Conv</span><span class="p">])</span>
<span class="n">WS</span> <span class="o">=</span> <span class="n">s</span><span class="o">.</span><span class="n">cache_read</span><span class="p">(</span><span class="n">W</span><span class="p">,</span> <span class="s2">&quot;shared&quot;</span><span class="p">,</span> <span class="p">[</span><span class="n">Conv</span><span class="p">])</span>
<span class="n">AF</span> <span class="o">=</span> <span class="n">s</span><span class="o">.</span><span class="n">cache_read</span><span class="p">(</span><span class="n">AS</span><span class="p">,</span> <span class="s2">&quot;wmma.matrix_a&quot;</span><span class="p">,</span> <span class="p">[</span><span class="n">Conv</span><span class="p">])</span>
<span class="n">WF</span> <span class="o">=</span> <span class="n">s</span><span class="o">.</span><span class="n">cache_read</span><span class="p">(</span><span class="n">WS</span><span class="p">,</span> <span class="s2">&quot;wmma.matrix_b&quot;</span><span class="p">,</span> <span class="p">[</span><span class="n">Conv</span><span class="p">])</span>
<span class="n">ConvF</span> <span class="o">=</span> <span class="n">s</span><span class="o">.</span><span class="n">cache_write</span><span class="p">(</span><span class="n">Conv</span><span class="p">,</span> <span class="s2">&quot;wmma.accumulator&quot;</span><span class="p">)</span>
</pre></div>
</div>
</div>
<div class="section" id="define-tensor-intrinsic">
<h2>Define Tensor Intrinsic<a class="headerlink" href="#define-tensor-intrinsic" title="永久链接至标题">¶</a></h2>
<p>In fact, TensorCore is a special hardware operation. So, we can just use tensorize
to replace a unit of computation with the TensorCore instruction. The first thing is
that we need to define tensor intrinsic.</p>
<p>There are four basic operation in TensorCore: <code class="code docutils literal notranslate"><span class="pre">fill_fragment</span></code>, <code class="code docutils literal notranslate"><span class="pre">load_matrix</span></code>,
<code class="code docutils literal notranslate"><span class="pre">mma_sync</span></code> and <code class="code docutils literal notranslate"><span class="pre">store_matrix</span></code>. Since <code class="code docutils literal notranslate"><span class="pre">fill_fragment</span></code> and <code class="code docutils literal notranslate"><span class="pre">mma_sync</span></code>
are both used in matrix multiplication, so we can just write following three intrinsics.</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">intrin_wmma_load_matrix</span><span class="p">(</span><span class="n">scope</span><span class="p">):</span>
    <span class="n">n</span> <span class="o">=</span> <span class="mi">16</span>
    <span class="n">A</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;A&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;float16&quot;</span><span class="p">)</span>
    <span class="n">BA</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">A</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">A</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">scope</span><span class="o">=</span><span class="s2">&quot;shared&quot;</span><span class="p">,</span> <span class="n">data_alignment</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">256</span><span class="p">)</span>
    <span class="n">C</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">((</span><span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">:</span> <span class="n">A</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">],</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;C&quot;</span><span class="p">)</span>
    <span class="n">BC</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">C</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">scope</span><span class="o">=</span><span class="n">scope</span><span class="p">,</span> <span class="n">data_alignment</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">256</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">intrin_func</span><span class="p">(</span><span class="n">ins</span><span class="p">,</span> <span class="n">outs</span><span class="p">):</span>
        <span class="n">ib</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ir_builder</span><span class="o">.</span><span class="n">create</span><span class="p">()</span>

        <span class="n">BA</span> <span class="o">=</span> <span class="n">ins</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">BC</span> <span class="o">=</span> <span class="n">outs</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">ib</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span>
            <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">call_intrin</span><span class="p">(</span>
                <span class="s2">&quot;handle&quot;</span><span class="p">,</span>
                <span class="s2">&quot;tir.tvm_load_matrix_sync&quot;</span><span class="p">,</span>
                <span class="n">BC</span><span class="o">.</span><span class="n">data</span><span class="p">,</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="n">BC</span><span class="o">.</span><span class="n">elem_offset</span> <span class="o">//</span> <span class="mi">256</span><span class="p">,</span>
                <span class="n">BA</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;r&quot;</span><span class="p">),</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="s2">&quot;row_major&quot;</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="p">)</span>
        <span class="k">return</span> <span class="n">ib</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">te</span><span class="o">.</span><span class="n">decl_tensor_intrin</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">op</span><span class="p">,</span> <span class="n">intrin_func</span><span class="p">,</span> <span class="n">binds</span><span class="o">=</span><span class="p">{</span><span class="n">A</span><span class="p">:</span> <span class="n">BA</span><span class="p">,</span> <span class="n">C</span><span class="p">:</span> <span class="n">BC</span><span class="p">})</span>


<span class="k">def</span> <span class="nf">intrin_wmma_gemm</span><span class="p">():</span>
    <span class="n">n</span> <span class="o">=</span> <span class="mi">16</span>
    <span class="n">A</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;A&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;float16&quot;</span><span class="p">)</span>
    <span class="n">B</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;B&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;float16&quot;</span><span class="p">)</span>
    <span class="n">k</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">reduce_axis</span><span class="p">((</span><span class="mi">0</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;k&quot;</span><span class="p">)</span>
    <span class="n">C</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">(</span>
        <span class="p">(</span><span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span>
        <span class="k">lambda</span> <span class="n">ii</span><span class="p">,</span> <span class="n">jj</span><span class="p">:</span> <span class="n">te</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">A</span><span class="p">[</span><span class="n">ii</span><span class="p">,</span> <span class="n">k</span><span class="p">]</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s2">&quot;float&quot;</span><span class="p">)</span> <span class="o">*</span> <span class="n">B</span><span class="p">[</span><span class="n">k</span><span class="p">,</span> <span class="n">jj</span><span class="p">]</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s2">&quot;float&quot;</span><span class="p">),</span> <span class="n">axis</span><span class="o">=</span><span class="n">k</span><span class="p">),</span>
        <span class="n">name</span><span class="o">=</span><span class="s2">&quot;C&quot;</span><span class="p">,</span>
    <span class="p">)</span>
    <span class="n">BA</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span>
        <span class="n">A</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">A</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;BA&quot;</span><span class="p">,</span> <span class="n">scope</span><span class="o">=</span><span class="s2">&quot;wmma.matrix_a&quot;</span><span class="p">,</span> <span class="n">data_alignment</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">256</span>
    <span class="p">)</span>
    <span class="n">BB</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span>
        <span class="n">B</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">B</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;BB&quot;</span><span class="p">,</span> <span class="n">scope</span><span class="o">=</span><span class="s2">&quot;wmma.matrix_b&quot;</span><span class="p">,</span> <span class="n">data_alignment</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">256</span>
    <span class="p">)</span>
    <span class="n">BC</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span>
        <span class="n">C</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">C</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;BC&quot;</span><span class="p">,</span> <span class="n">scope</span><span class="o">=</span><span class="s2">&quot;wmma.accumulator&quot;</span><span class="p">,</span> <span class="n">data_alignment</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">256</span>
    <span class="p">)</span>

    <span class="k">def</span> <span class="nf">intrin_func</span><span class="p">(</span><span class="n">ins</span><span class="p">,</span> <span class="n">outs</span><span class="p">):</span>
        <span class="n">BA</span><span class="p">,</span> <span class="n">BB</span> <span class="o">=</span> <span class="n">ins</span>
        <span class="p">(</span><span class="n">BC</span><span class="p">,)</span> <span class="o">=</span> <span class="n">outs</span>

        <span class="k">def</span> <span class="nf">init</span><span class="p">():</span>
            <span class="n">ib</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ir_builder</span><span class="o">.</span><span class="n">create</span><span class="p">()</span>
            <span class="n">ib</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span>
                <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">call_intrin</span><span class="p">(</span>
                    <span class="s2">&quot;handle&quot;</span><span class="p">,</span> <span class="s2">&quot;tir.tvm_fill_fragment&quot;</span><span class="p">,</span> <span class="n">BC</span><span class="o">.</span><span class="n">data</span><span class="p">,</span> <span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">,</span> <span class="n">BC</span><span class="o">.</span><span class="n">elem_offset</span> <span class="o">//</span> <span class="mi">256</span><span class="p">,</span> <span class="mf">0.0</span>
                <span class="p">)</span>
            <span class="p">)</span>
            <span class="k">return</span> <span class="n">ib</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

        <span class="k">def</span> <span class="nf">update</span><span class="p">():</span>
            <span class="n">ib</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ir_builder</span><span class="o">.</span><span class="n">create</span><span class="p">()</span>
            <span class="n">ib</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span>
                <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">call_intrin</span><span class="p">(</span>
                    <span class="s2">&quot;handle&quot;</span><span class="p">,</span>
                    <span class="s2">&quot;tir.tvm_mma_sync&quot;</span><span class="p">,</span>
                    <span class="n">BC</span><span class="o">.</span><span class="n">data</span><span class="p">,</span>
                    <span class="n">BC</span><span class="o">.</span><span class="n">elem_offset</span> <span class="o">//</span> <span class="mi">256</span><span class="p">,</span>
                    <span class="n">BA</span><span class="o">.</span><span class="n">data</span><span class="p">,</span>
                    <span class="n">BA</span><span class="o">.</span><span class="n">elem_offset</span> <span class="o">//</span> <span class="mi">256</span><span class="p">,</span>
                    <span class="n">BB</span><span class="o">.</span><span class="n">data</span><span class="p">,</span>
                    <span class="n">BB</span><span class="o">.</span><span class="n">elem_offset</span> <span class="o">//</span> <span class="mi">256</span><span class="p">,</span>
                    <span class="n">BC</span><span class="o">.</span><span class="n">data</span><span class="p">,</span>
                    <span class="n">BC</span><span class="o">.</span><span class="n">elem_offset</span> <span class="o">//</span> <span class="mi">256</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="p">)</span>
            <span class="k">return</span> <span class="n">ib</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

        <span class="k">return</span> <span class="n">update</span><span class="p">(),</span> <span class="n">init</span><span class="p">(),</span> <span class="n">update</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">te</span><span class="o">.</span><span class="n">decl_tensor_intrin</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">op</span><span class="p">,</span> <span class="n">intrin_func</span><span class="p">,</span> <span class="n">binds</span><span class="o">=</span><span class="p">{</span><span class="n">A</span><span class="p">:</span> <span class="n">BA</span><span class="p">,</span> <span class="n">B</span><span class="p">:</span> <span class="n">BB</span><span class="p">,</span> <span class="n">C</span><span class="p">:</span> <span class="n">BC</span><span class="p">})</span>


<span class="k">def</span> <span class="nf">intrin_wmma_store_matrix</span><span class="p">():</span>
    <span class="n">n</span> <span class="o">=</span> <span class="mi">16</span>
    <span class="n">A</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">placeholder</span><span class="p">((</span><span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;A&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;float32&quot;</span><span class="p">)</span>
    <span class="n">BA</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span>
        <span class="n">A</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">A</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">scope</span><span class="o">=</span><span class="s2">&quot;wmma.accumulator&quot;</span><span class="p">,</span> <span class="n">data_alignment</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">256</span>
    <span class="p">)</span>
    <span class="n">C</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">compute</span><span class="p">((</span><span class="n">n</span><span class="p">,</span> <span class="n">n</span><span class="p">),</span> <span class="k">lambda</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">:</span> <span class="n">A</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">],</span> <span class="n">name</span><span class="o">=</span><span class="s2">&quot;C&quot;</span><span class="p">)</span>
    <span class="n">BC</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">decl_buffer</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">C</span><span class="o">.</span><span class="n">dtype</span><span class="p">,</span> <span class="n">scope</span><span class="o">=</span><span class="s2">&quot;global&quot;</span><span class="p">,</span> <span class="n">data_alignment</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">offset_factor</span><span class="o">=</span><span class="mi">256</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">intrin_func</span><span class="p">(</span><span class="n">ins</span><span class="p">,</span> <span class="n">outs</span><span class="p">):</span>
        <span class="n">ib</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">ir_builder</span><span class="o">.</span><span class="n">create</span><span class="p">()</span>
        <span class="n">BA</span> <span class="o">=</span> <span class="n">ins</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">BC</span> <span class="o">=</span> <span class="n">outs</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">ib</span><span class="o">.</span><span class="n">emit</span><span class="p">(</span>
            <span class="n">tvm</span><span class="o">.</span><span class="n">tir</span><span class="o">.</span><span class="n">call_intrin</span><span class="p">(</span>
                <span class="s2">&quot;handle&quot;</span><span class="p">,</span>
                <span class="s2">&quot;tir.tvm_store_matrix_sync&quot;</span><span class="p">,</span>
                <span class="n">BA</span><span class="o">.</span><span class="n">data</span><span class="p">,</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="n">BA</span><span class="o">.</span><span class="n">elem_offset</span> <span class="o">//</span> <span class="mi">256</span><span class="p">,</span>
                <span class="n">BC</span><span class="o">.</span><span class="n">access_ptr</span><span class="p">(</span><span class="s2">&quot;w&quot;</span><span class="p">),</span>
                <span class="n">n</span><span class="p">,</span>
                <span class="s2">&quot;row_major&quot;</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="p">)</span>
        <span class="k">return</span> <span class="n">ib</span><span class="o">.</span><span class="n">get</span><span class="p">()</span>

    <span class="k">return</span> <span class="n">te</span><span class="o">.</span><span class="n">decl_tensor_intrin</span><span class="p">(</span><span class="n">C</span><span class="o">.</span><span class="n">op</span><span class="p">,</span> <span class="n">intrin_func</span><span class="p">,</span> <span class="n">binds</span><span class="o">=</span><span class="p">{</span><span class="n">A</span><span class="p">:</span> <span class="n">BA</span><span class="p">,</span> <span class="n">C</span><span class="p">:</span> <span class="n">BC</span><span class="p">})</span>
</pre></div>
</div>
</div>
<div class="section" id="scheduling-the-computation">
<h2>Scheduling the Computation<a class="headerlink" href="#scheduling-the-computation" title="永久链接至标题">¶</a></h2>
<p>To use TensorCores in TVM, we must schedule the computation into specific structure
to match the tensor intrinsic. The same as traditional GPU programs, we can also use
shared memory to boost the speed. If you have any questions about blocking and shared
memory, please refer <a class="reference internal" href="opt_conv_cuda.html#opt-conv-gpu"><span class="std std-ref">如何在GPU上优化卷积</span></a>.</p>
<p>In this example, each block contains 2x4 warps, and each warp calls 4x2 TensorCore
instructions. Thus, the output shape of each warp is 64x32 and each block outputs
128x128 titles. Due to the limit of shared memory space, we only load 2 blocks (2x128x128 tiles)
one time.</p>
<div class="admonition note">
<p class="admonition-title">注解</p>
<p><em>Warp-level Operation</em></p>
<p>Note that all TensorCore instructions are warp-level instructions, which means all 32 threads
in a warp should do this instruction simultaneously. Making theadIdx.x extent=32 is one of the
easiest way to solve this. Then We can bind threadIdx.x to any loops except those contain
TensorCore intrinsics directly or indirectly. Also note that it is not the unique solution.
The only thing we should do is to make sure all threads in a warp can call TensorCore at the same time.</p>
</div>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Define tiling sizes</span>
<span class="n">block_row_warps</span> <span class="o">=</span> <span class="mi">4</span>
<span class="n">block_col_warps</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">warp_row_tiles</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">warp_col_tiles</span> <span class="o">=</span> <span class="mi">4</span>
<span class="n">warp_size</span> <span class="o">=</span> <span class="mi">32</span>
<span class="n">chunk</span> <span class="o">=</span> <span class="mi">2</span>

<span class="n">block_x</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">thread_axis</span><span class="p">(</span><span class="s2">&quot;blockIdx.x&quot;</span><span class="p">)</span>
<span class="n">block_y</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">thread_axis</span><span class="p">(</span><span class="s2">&quot;blockIdx.y&quot;</span><span class="p">)</span>
<span class="n">block_z</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">thread_axis</span><span class="p">(</span><span class="s2">&quot;blockIdx.z&quot;</span><span class="p">)</span>
<span class="n">thread_x</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">thread_axis</span><span class="p">(</span><span class="s2">&quot;threadIdx.x&quot;</span><span class="p">)</span>
<span class="n">thread_y</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">thread_axis</span><span class="p">(</span><span class="s2">&quot;threadIdx.y&quot;</span><span class="p">)</span>
<span class="n">thread_z</span> <span class="o">=</span> <span class="n">te</span><span class="o">.</span><span class="n">thread_axis</span><span class="p">(</span><span class="s2">&quot;threadIdx.z&quot;</span><span class="p">)</span>

<span class="n">nc</span><span class="p">,</span> <span class="n">hc</span><span class="p">,</span> <span class="n">wc</span><span class="p">,</span> <span class="n">oc</span><span class="p">,</span> <span class="n">nnc</span><span class="p">,</span> <span class="n">ooc</span> <span class="o">=</span> <span class="n">Conv</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">axis</span>
<span class="n">block_k</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">fuse</span><span class="p">(</span><span class="n">hc</span><span class="p">,</span> <span class="n">wc</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">block_k</span><span class="p">,</span> <span class="n">block_z</span><span class="p">)</span>
<span class="n">nc</span><span class="p">,</span> <span class="n">nci</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">nc</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">warp_row_tiles</span><span class="p">)</span>
<span class="n">block_i</span><span class="p">,</span> <span class="n">nc</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">nc</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">block_row_warps</span><span class="p">)</span>
<span class="n">oc</span><span class="p">,</span> <span class="n">oci</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">oc</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">warp_col_tiles</span><span class="p">)</span>
<span class="n">block_j</span><span class="p">,</span> <span class="n">oc</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">oc</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">block_col_warps</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">reorder</span><span class="p">(</span><span class="n">block_k</span><span class="p">,</span> <span class="n">block_i</span><span class="p">,</span> <span class="n">block_j</span><span class="p">,</span> <span class="n">nc</span><span class="p">,</span> <span class="n">oc</span><span class="p">,</span> <span class="n">nci</span><span class="p">,</span> <span class="n">oci</span><span class="p">,</span> <span class="n">nnc</span><span class="p">,</span> <span class="n">ooc</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">block_i</span><span class="p">,</span> <span class="n">block_x</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">block_j</span><span class="p">,</span> <span class="n">block_y</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">nc</span><span class="p">,</span> <span class="n">thread_y</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">oc</span><span class="p">,</span> <span class="n">thread_z</span><span class="p">)</span>

<span class="c1"># Schedule local computation</span>
<span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">]</span><span class="o">.</span><span class="n">compute_at</span><span class="p">(</span><span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">],</span> <span class="n">oc</span><span class="p">)</span>
<span class="n">n</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">o</span><span class="p">,</span> <span class="n">nnf</span><span class="p">,</span> <span class="n">oof</span> <span class="o">=</span> <span class="n">ConvF</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">axis</span>
<span class="n">ko</span><span class="p">,</span> <span class="n">ki</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">ic</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">chunk</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">]</span><span class="o">.</span><span class="n">reorder</span><span class="p">(</span><span class="n">ko</span><span class="p">,</span> <span class="n">kh</span><span class="p">,</span> <span class="n">ki</span><span class="p">,</span> <span class="n">kw</span><span class="p">,</span> <span class="n">n</span><span class="p">,</span> <span class="n">o</span><span class="p">,</span> <span class="n">nnf</span><span class="p">,</span> <span class="n">oof</span><span class="p">,</span> <span class="n">ii</span><span class="p">)</span>

<span class="c1"># Move intermediate computation into each output compute tile</span>
<span class="n">s</span><span class="p">[</span><span class="n">AF</span><span class="p">]</span><span class="o">.</span><span class="n">compute_at</span><span class="p">(</span><span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">],</span> <span class="n">kw</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">WF</span><span class="p">]</span><span class="o">.</span><span class="n">compute_at</span><span class="p">(</span><span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">],</span> <span class="n">kw</span><span class="p">)</span>

<span class="c1"># Schedule for A&#39;s share memory</span>
<span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">compute_at</span><span class="p">(</span><span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">],</span> <span class="n">kh</span><span class="p">)</span>
<span class="n">n</span><span class="p">,</span> <span class="n">h</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">nn</span><span class="p">,</span> <span class="n">ii</span> <span class="o">=</span> <span class="n">AS</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">axis</span>
<span class="n">tx</span><span class="p">,</span> <span class="n">xo</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">n</span><span class="p">,</span> <span class="n">nparts</span><span class="o">=</span><span class="n">block_row_warps</span><span class="p">)</span>
<span class="n">ty</span><span class="p">,</span> <span class="n">yo</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">xo</span><span class="p">,</span> <span class="n">nparts</span><span class="o">=</span><span class="n">block_col_warps</span><span class="p">)</span>
<span class="n">t</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">fuse</span><span class="p">(</span><span class="n">nn</span><span class="p">,</span> <span class="n">ii</span><span class="p">)</span>
<span class="n">to</span><span class="p">,</span> <span class="n">ti</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">factor</span><span class="o">=</span><span class="n">warp_size</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">tx</span><span class="p">,</span> <span class="n">thread_y</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">ty</span><span class="p">,</span> <span class="n">thread_z</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">AS</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">ti</span><span class="p">,</span> <span class="n">thread_x</span><span class="p">)</span>

<span class="c1"># Schedule for W&#39;s share memory</span>
<span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">compute_at</span><span class="p">(</span><span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">],</span> <span class="n">kh</span><span class="p">)</span>
<span class="n">kh</span><span class="p">,</span> <span class="n">kw</span><span class="p">,</span> <span class="n">ic</span><span class="p">,</span> <span class="n">o</span><span class="p">,</span> <span class="n">ii</span><span class="p">,</span> <span class="n">oo</span> <span class="o">=</span> <span class="n">WS</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">axis</span>
<span class="n">tx</span><span class="p">,</span> <span class="n">xo</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">o</span><span class="p">,</span> <span class="n">nparts</span><span class="o">=</span><span class="n">block_row_warps</span><span class="p">)</span>
<span class="n">ty</span><span class="p">,</span> <span class="n">yo</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">xo</span><span class="p">,</span> <span class="n">nparts</span><span class="o">=</span><span class="n">block_col_warps</span><span class="p">)</span>
<span class="n">t</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">fuse</span><span class="p">(</span><span class="n">ii</span><span class="p">,</span> <span class="n">oo</span><span class="p">)</span>
<span class="n">to</span><span class="p">,</span> <span class="n">ti</span> <span class="o">=</span> <span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">nparts</span><span class="o">=</span><span class="n">warp_size</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">tx</span><span class="p">,</span> <span class="n">thread_y</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">ty</span><span class="p">,</span> <span class="n">thread_z</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="n">to</span><span class="p">,</span> <span class="n">thread_x</span><span class="p">)</span>
<span class="n">s</span><span class="p">[</span><span class="n">WS</span><span class="p">]</span><span class="o">.</span><span class="n">vectorize</span><span class="p">(</span><span class="n">ti</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">W</span><span class="p">,</span> <span class="n">Conv</span><span class="p">],</span> <span class="n">simple_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(A_1: handle, W_1: handle, Conv_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {Conv: Buffer(Conv_2: Pointer(float32), float32, [16, 14, 14, 32, 16, 16], []),
             A: Buffer(A_2: Pointer(float16), float16, [16, 14, 14, 16, 16, 16], []),
             W: Buffer(W_2: Pointer(float16), float16, [3, 3, 16, 32, 16, 16], [])}
  buffer_map = {A_1: A, W_1: W, Conv_1: Conv} {
  attr [IterVar(blockIdx.z: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;blockIdx.z&quot;)] &quot;thread_extent&quot; = 196;
  allocate(Conv.wmma.accumulator: Pointer(wmma.accumulator float32), float32, [2048]), storage_scope = wmma.accumulator;
  allocate(Apad.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
  allocate(W.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
  allocate(Apad.shared.wmma.matrix_a: Pointer(wmma.matrix_a float16), float16, [512]), storage_scope = wmma.matrix_a;
  allocate(W.shared.wmma.matrix_b: Pointer(wmma.matrix_b float16), float16, [1024]), storage_scope = wmma.matrix_b;
  attr [IterVar(blockIdx.x: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;blockIdx.x&quot;)] &quot;thread_extent&quot; = 2;
  attr [IterVar(blockIdx.y: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;blockIdx.y&quot;)] &quot;thread_extent&quot; = 4;
  attr [IterVar(threadIdx.y: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.y&quot;)] &quot;thread_extent&quot; = 4;
  attr [IterVar(threadIdx.z: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.z&quot;)] &quot;thread_extent&quot; = 2 {
    for (n.c.init: int32, 0, 2) {
      for (o.c.init: int32, 0, 4) {
        for (nn.c.init: int32, 0, 16) {
          for (oo.c.init: int32, 0, 16) {
            Conv.wmma.accumulator[((((n.c.init*1024) + (o.c.init*256)) + (nn.c.init*16)) + oo.c.init)] = 0f32
          }
        }
      }
    }
    for (ic.outer: int32, 0, 8) {
      for (kh: int32, 0, 3) {
        for (ax2: int32, 0, 3) {
          for (ax3: int32, 0, 2) {
            for (ax4.ax5.fused.outer: int32, 0, 8) {
              attr [IterVar(threadIdx.x: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.x&quot;)] &quot;thread_extent&quot; = 32;
              Apad.shared[((((((threadIdx.y*3072) + (threadIdx.z*1536)) + (ax2*512)) + (ax3*256)) + (ax4.ax5.fused.outer*32)) + threadIdx.x)] = @tir.if_then_else(((((1 &lt;= (floordiv(blockIdx.z, 14) + kh)) &amp;&amp; ((floordiv(blockIdx.z, 14) + kh) &lt; 15)) &amp;&amp; (1 &lt;= (ax2 + floormod(blockIdx.z, 14)))) &amp;&amp; ((ax2 + floormod(blockIdx.z, 14)) &lt; 15)), (float16*)A_2[(((((((((((blockIdx.x*6422528) + (threadIdx.y*1605632)) + (threadIdx.z*802816)) + (kh*57344)) + (blockIdx.z*4096)) + (ax2*4096)) + (ic.outer*512)) + (ax3*256)) + (ax4.ax5.fused.outer*32)) + threadIdx.x) - 61440)], 0f16, dtype=float16)
            }
          }
        }
        for (ax1: int32, 0, 3) {
          for (ax2_1: int32, 0, 2) {
            attr [IterVar(threadIdx.x, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.x&quot;)] &quot;thread_extent&quot; = 32;
            W.shared[ramp((((((ax1*4096) + (ax2_1*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)] = (float16x8*)W_2[ramp(((((((((kh*393216) + (ax1*131072)) + (ic.outer*16384)) + (ax2_1*8192)) + (blockIdx.y*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)]
          }
        }
        for (ic.inner: int32, 0, 2) {
          for (kw: int32, 0, 3) {
            for (ax0: int32, 0, 2) {
              for (ax4: int32, 0, 16) {
                for (ax5: int32, 0, 16) {
                  Apad.shared.wmma.matrix_a[(((ax0*256) + (ax4*16)) + ax5)] = (float16*)Apad.shared[((((((threadIdx.y*3072) + (ax0*1536)) + (kw*512)) + (ic.inner*256)) + (ax4*16)) + ax5)]
                }
              }
            }
            for (ax3_1: int32, 0, 4) {
              for (ax4_1: int32, 0, 16) {
                for (ax5_1: int32, 0, 16) {
                  W.shared.wmma.matrix_b[(((ax3_1*256) + (ax4_1*16)) + ax5_1)] = (float16*)W.shared[((((((kw*4096) + (ic.inner*2048)) + (threadIdx.z*1024)) + (ax3_1*256)) + (ax4_1*16)) + ax5_1)]
                }
              }
            }
            for (n.c: int32, 0, 2) {
              for (o.c: int32, 0, 4) {
                for (nn.c: int32, 0, 16) {
                  for (oo.c: int32, 0, 16) {
                    for (ii: int32, 0, 16) {
                      Conv.wmma.accumulator[((((n.c*1024) + (o.c*256)) + (nn.c*16)) + oo.c)] = ((float32*)Conv.wmma.accumulator[((((n.c*1024) + (o.c*256)) + (nn.c*16)) + oo.c)] + (cast(float32, (float16*)Apad.shared.wmma.matrix_a[(((n.c*256) + (nn.c*16)) + ii)])*cast(float32, (float16*)W.shared.wmma.matrix_b[(((o.c*256) + (ii*16)) + oo.c)])))
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
    for (n.inner: int32, 0, 2) {
      for (o.inner: int32, 0, 4) {
        for (nn: int32, 0, 16) {
          for (oo: int32, 0, 16) {
            Conv_2[(((((((((blockIdx.x*12845056) + (threadIdx.y*3211264)) + (n.inner*1605632)) + (blockIdx.z*8192)) + (blockIdx.y*2048)) + (threadIdx.z*1024)) + (o.inner*256)) + (nn*16)) + oo)] = (float32*)Conv.wmma.accumulator[((((n.inner*1024) + (o.inner*256)) + (nn*16)) + oo)]
          }
        }
      }
    }
  }
}
</pre></div>
</div>
</div>
<div class="section" id="lowering-computation-to-intrinsics">
<h2>Lowering Computation to Intrinsics<a class="headerlink" href="#lowering-computation-to-intrinsics" title="永久链接至标题">¶</a></h2>
<p>The last phase is to lower the computation loops down to TensorCore hardware intrinsics
by mapping the 2D convolution to tensor intrinsics</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">s</span><span class="p">[</span><span class="n">AF</span><span class="p">]</span><span class="o">.</span><span class="n">tensorize</span><span class="p">(</span><span class="n">AF</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">axis</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">],</span> <span class="n">intrin_wmma_load_matrix</span><span class="p">(</span><span class="s2">&quot;wmma.matrix_a&quot;</span><span class="p">))</span>
<span class="n">s</span><span class="p">[</span><span class="n">WF</span><span class="p">]</span><span class="o">.</span><span class="n">tensorize</span><span class="p">(</span><span class="n">WF</span><span class="o">.</span><span class="n">op</span><span class="o">.</span><span class="n">axis</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">],</span> <span class="n">intrin_wmma_load_matrix</span><span class="p">(</span><span class="s2">&quot;wmma.matrix_b&quot;</span><span class="p">))</span>
<span class="n">s</span><span class="p">[</span><span class="n">Conv</span><span class="p">]</span><span class="o">.</span><span class="n">tensorize</span><span class="p">(</span><span class="n">nnc</span><span class="p">,</span> <span class="n">intrin_wmma_store_matrix</span><span class="p">())</span>
<span class="n">s</span><span class="p">[</span><span class="n">ConvF</span><span class="p">]</span><span class="o">.</span><span class="n">tensorize</span><span class="p">(</span><span class="n">nnf</span><span class="p">,</span> <span class="n">intrin_wmma_gemm</span><span class="p">())</span>
<span class="nb">print</span><span class="p">(</span><span class="n">tvm</span><span class="o">.</span><span class="n">lower</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">W</span><span class="p">,</span> <span class="n">Conv</span><span class="p">],</span> <span class="n">simple_mode</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>primfn(A_1: handle, W_1: handle, Conv_1: handle) -&gt; ()
  attr = {&quot;from_legacy_te_schedule&quot;: True, &quot;global_symbol&quot;: &quot;main&quot;, &quot;tir.noalias&quot;: True}
  buffers = {Conv: Buffer(Conv_2: Pointer(float32), float32, [16, 14, 14, 32, 16, 16], []),
             A: Buffer(A_2: Pointer(float16), float16, [16, 14, 14, 16, 16, 16], []),
             W: Buffer(W_2: Pointer(float16), float16, [3, 3, 16, 32, 16, 16], [])}
  buffer_map = {A_1: A, W_1: W, Conv_1: Conv} {
  attr [IterVar(blockIdx.z: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;blockIdx.z&quot;)] &quot;thread_extent&quot; = 196;
  allocate(Conv.wmma.accumulator: Pointer(wmma.accumulator float32), float32, [2048]), storage_scope = wmma.accumulator;
  allocate(Apad.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
  allocate(W.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
  allocate(Apad.shared.wmma.matrix_a: Pointer(wmma.matrix_a float16), float16, [512]), storage_scope = wmma.matrix_a;
  allocate(W.shared.wmma.matrix_b: Pointer(wmma.matrix_b float16), float16, [1024]), storage_scope = wmma.matrix_b;
  attr [IterVar(blockIdx.x: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;blockIdx.x&quot;)] &quot;thread_extent&quot; = 2;
  attr [IterVar(blockIdx.y: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;blockIdx.y&quot;)] &quot;thread_extent&quot; = 4;
  attr [IterVar(threadIdx.y: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.y&quot;)] &quot;thread_extent&quot; = 4;
  attr [IterVar(threadIdx.z: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.z&quot;)] &quot;thread_extent&quot; = 2 {
    for (n.c.init: int32, 0, 2) {
      for (o.c.init: int32, 0, 4) {
        @tir.tvm_fill_fragment(Conv.wmma.accumulator, 16, 16, 16, ((n.c.init*4) + o.c.init), 0f32, dtype=handle)
      }
    }
    for (ic.outer: int32, 0, 8) {
      for (kh: int32, 0, 3) {
        for (ax2: int32, 0, 3) {
          for (ax3: int32, 0, 2) {
            for (ax4.ax5.fused.outer: int32, 0, 8) {
              attr [IterVar(threadIdx.x: int32, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.x&quot;)] &quot;thread_extent&quot; = 32;
              Apad.shared[((((((threadIdx.y*3072) + (threadIdx.z*1536)) + (ax2*512)) + (ax3*256)) + (ax4.ax5.fused.outer*32)) + threadIdx.x)] = @tir.if_then_else(((((1 &lt;= (floordiv(blockIdx.z, 14) + kh)) &amp;&amp; ((floordiv(blockIdx.z, 14) + kh) &lt; 15)) &amp;&amp; (1 &lt;= (ax2 + floormod(blockIdx.z, 14)))) &amp;&amp; ((ax2 + floormod(blockIdx.z, 14)) &lt; 15)), (float16*)A_2[(((((((((((blockIdx.x*6422528) + (threadIdx.y*1605632)) + (threadIdx.z*802816)) + (kh*57344)) + (blockIdx.z*4096)) + (ax2*4096)) + (ic.outer*512)) + (ax3*256)) + (ax4.ax5.fused.outer*32)) + threadIdx.x) - 61440)], 0f16, dtype=float16)
            }
          }
        }
        for (ax1: int32, 0, 3) {
          for (ax2_1: int32, 0, 2) {
            attr [IterVar(threadIdx.x, (nullptr), &quot;ThreadIndex&quot;, &quot;threadIdx.x&quot;)] &quot;thread_extent&quot; = 32;
            W.shared[ramp((((((ax1*4096) + (ax2_1*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)] = (float16x8*)W_2[ramp(((((((((kh*393216) + (ax1*131072)) + (ic.outer*16384)) + (ax2_1*8192)) + (blockIdx.y*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)]
          }
        }
        for (ic.inner: int32, 0, 2) {
          for (kw: int32, 0, 3) {
            for (ax0: int32, 0, 2) {
              @tir.tvm_load_matrix_sync(Apad.shared.wmma.matrix_a, 16, 16, 16, ax0, @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float16), Apad.shared, ((((threadIdx.y*3072) + (ax0*1536)) + (kw*512)) + (ic.inner*256)), 256, 1, dtype=handle), 16, &quot;row_major&quot;, dtype=handle)
            }
            for (ax3_1: int32, 0, 4) {
              @tir.tvm_load_matrix_sync(W.shared.wmma.matrix_b, 16, 16, 16, ax3_1, @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float16), W.shared, ((((kw*4096) + (ic.inner*2048)) + (threadIdx.z*1024)) + (ax3_1*256)), 256, 1, dtype=handle), 16, &quot;row_major&quot;, dtype=handle)
            }
            for (n.c: int32, 0, 2) {
              for (o.c: int32, 0, 4) {
                @tir.tvm_mma_sync(Conv.wmma.accumulator, ((n.c*4) + o.c), Apad.shared.wmma.matrix_a, n.c, W.shared.wmma.matrix_b, o.c, Conv.wmma.accumulator, ((n.c*4) + o.c), dtype=handle)
              }
            }
          }
        }
      }
    }
    for (n.inner: int32, 0, 2) {
      for (o.inner: int32, 0, 4) {
        @tir.tvm_store_matrix_sync(Conv.wmma.accumulator, 16, 16, 16, ((n.inner*4) + o.inner), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), Conv_2, (((((((blockIdx.x*12845056) + (threadIdx.y*3211264)) + (n.inner*1605632)) + (blockIdx.z*8192)) + (blockIdx.y*2048)) + (threadIdx.z*1024)) + (o.inner*256)), 256, 2, dtype=handle), 16, &quot;row_major&quot;, dtype=handle)
      }
    }
  }
}
</pre></div>
</div>
</div>
<div class="section" id="generate-cuda-kernel">
<h2>Generate CUDA Kernel<a class="headerlink" href="#generate-cuda-kernel" title="永久链接至标题">¶</a></h2>
<p>Finally we use TVM to generate and compile the CUDA kernel, and evaluate the latency of convolution.
Since TensorCores are only supported in NVIDIA GPU with Compute Capability 7.0 or higher, it may not
be able to run on our build server</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="n">dev</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">cuda</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="k">if</span> <span class="n">nvcc</span><span class="o">.</span><span class="n">have_tensorcore</span><span class="p">(</span><span class="n">dev</span><span class="o">.</span><span class="n">compute_version</span><span class="p">):</span>
    <span class="k">with</span> <span class="n">tvm</span><span class="o">.</span><span class="n">transform</span><span class="o">.</span><span class="n">PassContext</span><span class="p">(</span><span class="n">config</span><span class="o">=</span><span class="p">{</span><span class="s2">&quot;tir.UnrollLoop&quot;</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;auto_max_step&quot;</span><span class="p">:</span> <span class="mi">16</span><span class="p">}}):</span>
        <span class="n">func</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">build</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="p">[</span><span class="n">A</span><span class="p">,</span> <span class="n">W</span><span class="p">,</span> <span class="n">Conv</span><span class="p">],</span> <span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
    <span class="n">a_np</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="n">data_shape</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">A</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
    <span class="n">w_np</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">size</span><span class="o">=</span><span class="n">kernel_shape</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="n">W</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
    <span class="n">a</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">a_np</span><span class="p">,</span> <span class="n">dev</span><span class="p">)</span>
    <span class="n">w</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">w_np</span><span class="p">,</span> <span class="n">dev</span><span class="p">)</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">tvm</span><span class="o">.</span><span class="n">nd</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">output_shape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">Conv</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> <span class="n">dev</span><span class="p">)</span>
    <span class="n">evaluator</span> <span class="o">=</span> <span class="n">func</span><span class="o">.</span><span class="n">time_evaluator</span><span class="p">(</span><span class="n">func</span><span class="o">.</span><span class="n">entry_name</span><span class="p">,</span> <span class="n">dev</span><span class="p">,</span> <span class="n">number</span><span class="o">=</span><span class="mi">10</span><span class="p">)</span>
    <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;conv2d with tensor core: </span><span class="si">%f</span><span class="s2"> ms&quot;</span> <span class="o">%</span> <span class="p">(</span><span class="n">evaluator</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">c</span><span class="p">)</span><span class="o">.</span><span class="n">mean</span> <span class="o">*</span> <span class="mf">1e3</span><span class="p">))</span>
</pre></div>
</div>
<p class="sphx-glr-script-out">输出:</p>
<div class="sphx-glr-script-out highlight-none notranslate"><div class="highlight"><pre><span></span>conv2d with tensor core: 4.579600 ms
</pre></div>
</div>
</div>
<div class="section" id="summary">
<h2>总结<a class="headerlink" href="#summary" title="永久链接至标题">¶</a></h2>
<p>This tutorial demonstrates how TVM scheduling primitives can be used to
call TensorCores on specific GPUs.</p>
<div class="sphx-glr-footer class sphx-glr-footer-example docutils container" id="sphx-glr-download-how-to-optimize-operators-opt-conv-tensorcore-py">
<div class="sphx-glr-download docutils container">
<p><a class="reference download internal" download="" href="../../_downloads/7372db5919b5619bc34fde3434862bca/opt_conv_tensorcore.py"><code class="xref download docutils literal notranslate"><span class="pre">Python</span> <span class="pre">源码下载:</span> <span class="pre">opt_conv_tensorcore.py</span></code></a></p>
</div>
<div class="sphx-glr-download docutils container">
<p><a class="reference download internal" download="" href="../../_downloads/7455981870c23c8c76482dedf33d8a42/opt_conv_tensorcore.ipynb"><code class="xref download docutils literal notranslate"><span class="pre">Jupyter</span> <span class="pre">notebook</span> <span class="pre">下载:</span> <span class="pre">opt_conv_tensorcore.ipynb</span></code></a></p>
</div>
</div>
<p class="sphx-glr-signature"><a class="reference external" href="https://sphinx-gallery.github.io">Gallery generated by Sphinx-Gallery</a></p>
</div>
</div>


           </div>
           
          </div>
          

<footer>

    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="../tune_with_autotvm/index.html" class="btn btn-neutral float-right" title="Auto-Tune with Templates and AutoTVM" accesskey="n" rel="next">下一个 <span class="fa fa-arrow-circle-right"></span></a>
      
      
        <a href="opt_conv_cuda.html" class="btn btn-neutral float-left" title="如何在GPU上优化卷积" accesskey="p" rel="prev"><span class="fa fa-arrow-circle-left"></span> 上一个</a>
      
    </div>

<div id="button" class="backtop"><img src="../../_static//img/right.svg" alt="backtop"/> </div>
<section class="footerSec">
    <div class="footerHeader">
      <ul class="d-flex align-md-items-center justify-content-between flex-column flex-md-row">
        <li class="copywrite d-flex align-items-center">
          <h5 id="copy-right-info">© 2020 Apache Software Foundation | All right reserved</h5>
        </li>
      </ul>

    </div>

    <ul>
      <li class="footernote">Copyright © 2020 The Apache Software Foundation. Apache TVM, Apache, the Apache feather, and the Apache TVM project logo are either trademarks or registered trademarks of the Apache Software Foundation.</li>
    </ul>

</section>
</footer>
        </div>
      </div>

    </section>

  </div>
  

    <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.12.9/umd/popper.min.js" integrity="sha384-ApNbgh9B+Y1QKtv3Rn7W3mgPxhU9K/ScQsAP7hUibX39j7fakFPskvXusvfa0b4Q" crossorigin="anonymous"></script>
    <script src="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0/js/bootstrap.min.js" integrity="sha384-JZR6Spejh4U02d8jOt6vLEHfe/JQGiRRSQQxSfFWpi1MquVdAyjUar5+76PVCmYl" crossorigin="anonymous"></script>

  </body>
  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
    <!-- Theme Analytics -->
    <script>
    (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){
      (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o),
      m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m)
    })(window,document,'script','https://www.google-analytics.com/analytics.js','ga');

    ga('create', 'UA-75982049-2', 'auto');
    ga('send', 'pageview');
    </script>

    
   

</body>
</html>